package com.caochenlei.tree.rbt;

import java.util.LinkedList;
import java.util.Queue;

public class RedBlackTree<Key extends Comparable<Key>, Value> {
    public class Node {                         //结点类
        public Key key;                         //存储键
        public Value value;                     //存储值
        public Node parent;                     //指向父亲结点
        public Node left;                       //指向左子结点
        public Node right;                      //指向右子结点
        public boolean color;                   //指向结点颜色

        public Node(Key key, Value value) {
            this.key = key;
            this.value = value;
        }

        @Override
        public String toString() {
            return "Node{" + "key=" + key + ", value=" + value + ", color=" + color + "}";
        }
    }

    private Node root;                          //根结点
    private int size;                           //结点数

    private static final boolean RED = false;   //红色结点
    private static final boolean BLACK = true;  //黑色结点

    //获取当前树的根结点
    public Node getRoot() {
        return root;
    }

    //获取当前树的结点数
    public int size() {
        return size;
    }

    //判断当前树是否为空
    public boolean isEmpty() {
        return size == 0;
    }

    //获取当前结点的父亲结点
    private Node parentOf(Node node) {
        return node != null ? node.parent : null;
    }

    //获取当前结点的左子结点
    private Node leftOf(Node node) {
        return node != null ? node.left : null;
    }

    //获取当前结点的右子结点
    private Node rightOf(Node node) {
        return node != null ? node.right : null;
    }

    //获取当前结点颜色
    private boolean colorOf(Node node) {
        return node != null ? node.color : BLACK;
    }

    //设置当前结点颜色
    private void setColor(Node node, boolean color) {
        if (node != null) {
            node.color = color;
        }
    }

    //左旋操作
    private void rotateLeft(Node h) {
        if (h != null) {                    //判断当前结点h是否为null
            Node x = h.right;               //获取当前结点h的右子结点x
            h.right = x.left;               //让h的右子结点指向x的左子结点
            if (x.left != null)             //判断x的左子结点是否为null
                x.left.parent = h;          //让x的左子结点的父结点指向h
            x.parent = h.parent;            //让x的父结点指向h的父结点
            if (h.parent == null)           //判断h的父结点是否为null
                root = x;                   //如果是，说明是根结点
            else if (h.parent.left == h)    //判断当前h子树是否为父结点的左子树
                h.parent.left = x;          //如果是，则修正他为新的x子树
            else                            //判断当前h子树是否为父结点的右子树
                h.parent.right = x;         //如果是，则修正他为新的x子树
            x.left = h;                     //让x的左子结点指向h结点
            h.parent = x;                   //让h结点的父结点指向x结点
        }
    }

    //右旋操作
    private void rotateRight(Node h) {
        if (h != null) {                    //判断当前结点h是否为null
            Node x = h.left;                //获取当前结点h的左子结点x
            h.left = x.right;               //让h的左子结点指向x的右子结点
            if (x.right != null)            //判断x的右子结点是否为null
                x.right.parent = h;         //让x的右子结点的父结点指向h
            x.parent = h.parent;            //让x的父结点指向h的父结点
            if (h.parent == null)           //判断h的父结点是否为null
                root = x;                   //如果是，说明是根结点
            else if (h.parent.right == h)   //判断当前h子树是否为父结点的右子树
                h.parent.right = x;         //如果是，则修正他为新的x子树
            else                            //判断当前h子树是否为父结点的左子树
                h.parent.left = x;          //如果是，则修正他为新的x子树
            x.right = h;                    //让x的右子结点指向h结点
            h.parent = x;                   //让h结点的父结点指向x结点
        }
    }

    //添加操作
    public void put(Key key, Value value) {
        //获取根结点
        Node t = root;
        //判断根结点
        if (t == null) {
            root = new Node(key, value);
            root.color = BLACK;
            size++;
            return;
        }
        //找指定位置
        Node parent;
        int cmp;
        do {
            parent = t;
            cmp = key.compareTo(t.key);
            if (cmp < 0)
                t = t.left;
            else if (cmp > 0)
                t = t.right;
            else
                t.value = value;
        } while (t != null);
        //添加新结点
        Node e = new Node(key, value);
        e.parent = parent;
        if (cmp < 0)
            parent.left = e;
        else
            parent.right = e;
        //添加后修正
        fixAfterPut(e);
        //添加后加一
        size++;
    }

    //修正操作
    private void fixAfterPut(Node x) {
        x.color = RED;//每次添加必须是红结点

        //因为root没有父结点，所以不能让x上移循环到root
        while (x != null && x != root && x.parent.color == RED) {
            Node grandpa = parentOf(parentOf(x));
            if (parentOf(x) == leftOf(grandpa)) {
                Node uncle = rightOf(grandpa);
                if (colorOf(uncle) == RED) {//这个是判断的：4、新增一个结点，与4结点合并，此时分裂
                    setColor(parentOf(x), BLACK);
                    setColor(uncle, BLACK);
                    setColor(grandpa, RED);
                    x = grandpa;//因为x结点、x的父结点、x的爷爷结点都修正过了，该让x的爷爷结点往上继续修正了
                } else {//这里是判断的：3、新增一个结点，与3结点合并，直接合并
                    if (x == rightOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateLeft(x);
                    }
                    setColor(parentOf(x), BLACK);
                    setColor(grandpa, RED);
                    rotateRight(grandpa);
                }
            } else {
                Node uncle = leftOf(grandpa);
                if (colorOf(uncle) == RED) {//这个是判断的：4、新增一个结点，与4结点合并，此时分裂
                    setColor(parentOf(x), BLACK);
                    setColor(uncle, BLACK);
                    setColor(grandpa, RED);
                    x = grandpa;//因为x结点、x的父结点、x的爷爷结点都修正过了，该让x的爷爷结点往上继续修正了
                } else {//这里是判断的：3、新增一个结点，与3结点合并，直接合并
                    if (x == leftOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateRight(x);
                    }
                    setColor(parentOf(x), BLACK);
                    setColor(grandpa, RED);
                    rotateLeft(grandpa);
                }
            }
        }

        root.color = BLACK;//根结点必须是黑结点
    }

    //层序遍历当前树
    public void layerErgodic() {
        layerErgodic(root);
    }

    //层序遍历指定树
    private void layerErgodic(Node x) {
        //创建一个队列
        Queue<Node> nodes = new LinkedList<>();
        //加入指定节点
        nodes.offer(x);
        //循环弹出遍历
        while (!nodes.isEmpty()) {
            //从队列中弹出一个结点，输出当前结点的信息
            Node n = nodes.poll();
            System.out.println(n);
            //判断当前结点还有没有左子结点，如果有，则放入到nodes中
            if (n.left != null) {
                nodes.offer(n.left);
            }
            //判断当前结点还有没有右子结点，如果有，则放入到nodes中
            if (n.right != null) {
                nodes.offer(n.right);
            }
        }
    }

    //获取树指定结点
    public Node getNode(Key key) {
        Node p = root;
        while (p != null) {
            int cmp = key.compareTo(p.key);
            if (cmp < 0)
                p = p.left;
            else if (cmp > 0)
                p = p.right;
            else
                return p;
        }
        return null;
    }

    //获取树最小结点
    public Node getMinNode() {
        Node p = root;
        if (p != null)
            while (p.left != null)
                p = p.left;
        return p;
    }

    //获取树最大结点
    public Node getMaxNode() {
        Node p = root;
        if (p != null)
            while (p.right != null)
                p = p.right;
        return p;
    }

    //获取指定结点的前驱结点
    public Node predecessor(Node t) {
        if (t == null)//第一种情况
            return null;
        else if (t.left != null) {//第二种情况
            Node p = t.left;
            while (p.right != null)
                p = p.right;
            return p;
        } else {//第三种情况
            Node p = t.parent;
            Node ch = t;
            while (p != null && ch == p.left) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    //获取指定结点的后继结点
    public Node successor(Node t) {
        if (t == null)//第一种情况
            return null;
        else if (t.right != null) {//第二种情况
            Node p = t.right;
            while (p.left != null)
                p = p.left;
            return p;
        } else {//第三种情况
            Node p = t.parent;
            Node ch = t;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    //删除操作
    public Value delete(Key key) {
        Node p = getNode(key);
        if (p == null)
            return null;
        Value oldValue = p.value;
        deleteNode(p);
        return oldValue;
    }

    //真正删除
    private void deleteNode(Node p) {
        size--;

        //4、要删除的结点有2个孩子，肯定有左孩子，肯定有右孩子
        if (p.left != null && p.right != null) {
            //找到前驱结点，然后替换要删除结点的键值
            Node s = predecessor(p);
            p.key = s.key;
            p.value = s.value;
            //替换完成以后，指向前驱结点为删除做准备
            p = s;
        }

        //3、要删除的结点有1个孩子，可能有左孩子，可能有右孩子
        Node replacement = (p.left != null ? p.left : p.right);
        if (replacement != null) {
            //让左右子结点连接到p的父结点上
            replacement.parent = p.parent;
            if (p.parent == null)
                root = replacement;
            else if (p == p.parent.left)
                p.parent.left = replacement;
            else
                p.parent.right = replacement;

            //释放当前p结点所有指向等待回收
            p.left = p.right = p.parent = null;

            //如果当前删除的是黑结点需修正
            if (p.color == BLACK)
                fixAfterDelete(replacement);
        }
        //2、要删除的结点是树根结点，此时直接删除即可。（树根结点永远黑色，删除完了就完了，还保持啥平衡）
        else if (p.parent == null) {
            root = null;
        }
        //1、要删除的结点是叶子结点，此时直接删除即可。（如果删除的是黑色结点，我们要先调整，然后再删除）
        else {
            //先调整
            if (p.color == BLACK)
                fixAfterDelete(p);

            //再删除
            if (p.parent != null) {
                if (p == p.parent.left)
                    p.parent.left = null;
                else if (p == p.parent.right)
                    p.parent.right = null;
                p.parent = null;
            }
        }
    }

    //修正操作
    private void fixAfterDelete(Node x) {
        //那么在此之前要处理的必定是  x != root && colorOf(x) == BLACK
        while (x != root && colorOf(x) == BLACK) {
            //x是左孩子的情况
            if (x == leftOf(parentOf(x))) {
                //获取当前的兄弟结点
                Node brother = rightOf(parentOf(x));

                //找到真正的兄弟结点
                if (colorOf(brother) == RED) {
                    setColor(brother, BLACK);
                    setColor(parentOf(x), RED);
                    rotateLeft(parentOf(x));
                    brother = rightOf(parentOf(x));
                }

                //这个兄弟结点帮不了
                if (colorOf(leftOf(brother)) == BLACK && colorOf(rightOf(brother)) == BLACK) {
                    setColor(brother, RED);
                    x = parentOf(x);
                }
                //这个兄弟结点帮得了
                else {
                    //判断当前结构是不是需要调整，找到真正的兄弟结点
                    if (colorOf(rightOf(brother)) == BLACK) {
                        setColor(leftOf(brother), BLACK);
                        setColor(brother, RED);
                        rotateRight(brother);
                        brother = rightOf(parentOf(x));
                    }
                    //让老父亲去顶替被删除结点，让亲兄弟去顶替老父亲
                    setColor(brother, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(rightOf(brother), BLACK);
                    rotateLeft(parentOf(x));
                    //这种情况，调整一次即可，x=root代表跳出当前循环
                    x = root;
                }
            }
            //x是右孩子的情况
            else {
                //获取当前的兄弟结点
                Node brother = leftOf(parentOf(x));

                //找到真正的兄弟结点
                if (colorOf(brother) == RED) {
                    setColor(brother, BLACK);
                    setColor(parentOf(x), RED);
                    rotateRight(parentOf(x));
                    brother = leftOf(parentOf(x));
                }

                //这个兄弟结点帮不了
                if (colorOf(rightOf(brother)) == BLACK && colorOf(leftOf(brother)) == BLACK) {
                    setColor(brother, RED);
                    x = parentOf(x);
                }
                //这个兄弟结点帮得了
                else {
                    //判断当前结构是不是需要调整，找到真正的兄弟结点
                    if (colorOf(leftOf(brother)) == BLACK) {
                        setColor(rightOf(brother), BLACK);
                        setColor(brother, RED);
                        rotateLeft(brother);
                        brother = leftOf(parentOf(x));
                    }
                    //让老父亲去顶替被删除结点，让亲兄弟去顶替老父亲
                    setColor(brother, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(leftOf(brother), BLACK);
                    rotateRight(parentOf(x));
                    //这种情况，调整一次即可，x=root代表跳出当前循环
                    x = root;
                }
            }
        }

        //如果当前结点是红色，那么立即修改为黑色
        //还有一种情况就是为了保证根结点root始终是黑色，因为replacement最后有可能成为根结点
        setColor(x, BLACK);
    }
}
